from datetime import datetime, timedelta

import chinese_calendar as cc  # 用于检查中国法定节假日
import matplotlib
import matplotlib.pyplot as plt
import pandas as pd
import pywencai

# 设置中文字体
matplotlib.rcParams['font.sans-serif'] = ['SimHei']  # 使用SimHei字体
matplotlib.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题


# 获取从结束日期向前推的n个工作日（剔除中国法定节假日和周末）
def get_last_n_business_days(n, end_date):
    business_days = []
    current_date = pd.to_datetime(end_date)

    while len(business_days) < n:
        if current_date.weekday() < 5 and not cc.is_holiday(current_date):  # 排除周末和法定节假日
            business_days.append(current_date.strftime('%Y-%m-%d'))
        current_date -= timedelta(days=1)

    return business_days[::-1]  # 按时间升序排列


# 生成查询字符串
def get_queries(start_date, n_queries, days_per_query):
    queries = []
    end_date = pd.to_datetime(start_date)

    for _ in range(n_queries):
        business_days = get_last_n_business_days(days_per_query, end_date)
        query = f"从{business_days[0]}到{business_days[-1]}涨幅排名前10，创业板不含ST"
        queries.append(query)
        end_date = pd.to_datetime(business_days[-2])  # 向前移动一个工作日

    return queries


# 执行查询
def do_query(query):
    return pywencai.get(query=query)


# 绘制折线图
def plot_line_chart(all_data):
    # 转换为DataFrame
    all_data_df = pd.DataFrame(all_data)

    # 日期列转换为datetime格式并按升序排列
    all_data_df['日期'] = pd.to_datetime(all_data_df['日期'])
    all_data_df = all_data_df.sort_values(by='日期', ascending=True)

    # 创建图表
    plt.figure(figsize=(12, 6))

    # 绘制每只股票的折线
    for stock in all_data_df['股票名称'].unique():
        stock_data = all_data_df[all_data_df['股票名称'] == stock]
        stock_data = stock_data.sort_values(by='日期', ascending=True)
        plt.plot(stock_data['日期'], stock_data['近十日涨幅'], marker='o', label=stock)

        # 标注最后一天的数据点
        last_point = stock_data.iloc[-1]
        plt.text(last_point['日期'], last_point['近十日涨幅'], last_point['股票名称'],
                 fontsize=10, verticalalignment='bottom', horizontalalignment='left')

    # 设置工作日为横坐标
    workdays = pd.bdate_range(start=all_data_df['日期'].min(), end=all_data_df['日期'].max())
    plt.xticks(workdays, workdays.strftime('%Y-%m-%d'), rotation=45)

    # 图表标题和标签
    plt.title("十日涨幅折线图", fontsize=16)
    plt.xlabel("日期", fontsize=12)
    plt.ylabel("近十日涨幅 (%)", fontsize=12)
    plt.legend(loc='upper left', fontsize=10, bbox_to_anchor=(1, 1))  # 图例放在图表外
    plt.grid(True, linestyle='--', alpha=0.6)
    plt.tight_layout()
    plt.show()


if __name__ == '__main__':
    start_date = datetime.today().strftime('%Y-%m-%d')
    n_queries = 6  # 查询次数
    days_per_query = 10  # 每次查询的工作日数量

    # 获取查询列表
    queries = get_queries(start_date, n_queries, days_per_query)
    all_data = []

    # 执行查询并处理数据
    for query in queries:
        print(query)
        res = do_query(query)
        # 假设结果是DataFrame，第2列是股票名称，第6列是涨幅
        stock_names = res.iloc[:, 1]
        changes = res.iloc[:, 5]
        end_date = query.split('到')[1].split('涨幅')[0]
        for stock, change in zip(stock_names, changes):
            all_data.append({'股票名称': stock, '近十日涨幅': change, '日期': end_date})

    # 绘制图表
    plot_line_chart(all_data)
